


source /root/miniconda3/etc/profile.d/conda.sh
conda activate ws
atk_index=$1
rm -rf /etc/Ascend/ascend_cann_install.info

mkdir /home/ws/${atk_index}
mkdir /home/ws/${atk_index}/cann-package
mkdir /home/ws/${atk_index}/cann-ops-adv-dev
mkdir /home/ws/${atk_index}/result
mkdir /home/ws/${atk_index}/result/log

# 提取以tests/DT/和src/transformer/开头的路径的下一级目录名
cd /home/ws/${atk_index}
grep -E '^(tests/DT/|src/transformer/attention/)' pr_filelist.txt | \
awk -F'/' '{
    if ($1 == "tests" && $2 == "DT") print $3;
    else if ($1 == "src" && $2 == "transformer" && $3 == "attention") print $4
}' | \
sort -u >> case_name.txt

# 逐行读取 case_name.txt
while IFS= read -r case_name; do
    # 跳过空行
    if [ -z "$case_name" ]; then
        continue
    fi
    # 在 /home/ws/atk2case.csv 中查找匹配第一列的行，提取第二列
    awk -F, -v case_var="$case_name" '$1 == case_var {print $2}' /home/ws/atk2case.csv >> compile_name.txt
done < case_name.txt

# 逐行读取 compile_name.txt
while IFS= read -r target_case; do
    # 空行过滤
    [[ -z "$target_case" ]] && continue
    # 使用awk进行精确匹配（规避gawk关键字冲突）
    awk -F',' -v tcase="$target_case" '
        $1 == tcase {
            # 输出第2列到最后一列(兼容不定长字段)
            for(i=2; i<=NF; i++) {
                printf "%s%s", $i, (i==NF ? "\n" : ",")
            }
        }
    ' /home/ws/case2atk.csv >> "atk_name.txt"
done < compile_name.txt

perl -i -pe 's/,/\n/g' atk_name.txt

# cp case_name.txt compile_name.txt
sed -i '/fused_infer_attention_score/ {
    s//incre_flash_attention/;   
    a\prompt_flash_attention      
}' compile_name.txt
sort -u -o compile_name.txt compile_name.txt

# 生成参数列表
params=$(awk '{printf "-n %s ", $0}' compile_name.txt)

# 组合完整命令
command="bash build.sh ${params}--disable-check-compatible"

if [ $(wc -l < "/home/ws/${atk_index}/case_name.txt") -eq 0 ]; then
    echo "SKIP ATK-test"
    cat /home/ws/pass.log >> /home/ws/${atk_index}/result/log/atk.log
else
    # 安装cann
    install_path=/home/ws/${atk_index}/cann-package
    chmod -R 755 ${install_path}
    export TMPDIR=/dev/shm
    echo ">>>>>>>>>> start install cann <<<<<<<<<<"
    /home/ws/Ascend-cann-toolkit_*_linux-aarch64.run --full --quiet --install-path=$install_path
    /home/ws/Ascend-cann-kernels-*_linux-aarch64.run --install --quiet --install-path=$install_path
    echo ">>>>>>>>>> end install cann <<<<<<<<<<"
    # 编译
    echo ">>>>>>>>>> start compile <<<<<<<<<<"
    source /home/ws/${atk_index}/cann-package/ascend-toolkit/set_env.sh
    cd /home/ws/${atk_index}/cann-ops-adv-dev
    eval "$command"
    /home/ws/${atk_index}/cann-ops-adv-dev/output/CANN-custom_ops-*linux.aarch64.run --quiet
    export LD_LIBRARY_PATH=/home/ws/${atk_index}/cann-package/ascend-toolkit/latest/opp/vendors/customize/op_api/lib/:${LD_LIBRARY_PATH}
    # 执行
    source /home/ws/${atk_index}/cann-package/ascend-toolkit/set_env.sh
    echo ">>>>>>>>>> start run atk <<<<<<<<<<"
    for case_name in $(cat /home/ws/${atk_index}/atk_name.txt)
    do
        cd /home/ws/${atk_index}/cann-ops-adv-dev/tests/DT/${case_name}
        atk task -c ${case_name}.json -n node.yaml -p ${case_name}.py
        mkdir /home/ws/${atk_index}/result/${case_name}
        cp -r /home/ws/${atk_index}/cann-ops-adv-dev/tests/DT/${case_name}/atk_output/${case_name}_*/* /home/ws/${atk_index}/result/${case_name}
        # 判断是否通过atk
        has_pass=$(grep -i -m1 -E "Pass" /home/ws/${atk_index}/result/${case_name}/log/atk.log | wc -l)
        # # [ $has_pass -gt 0 ] && echo true || echo false
        if [ $has_pass -gt 0 ];then
            echo "${case_name} ATK-test pass"
        else
            echo "${case_name} ATK-test faild"
        fi
        cat /home/ws/${atk_index}/result/${case_name}/log/atk.log >> /home/ws/${atk_index}/result/log/atk.log
    done
fi